#include "stream_reassembler.hh"

#include <utility>

// Dummy implementation of a stream reassembler.

// For Lab 1, please replace with a real implementation that passes the
// automated checks run by `make check_lab1`.

// You will need to add private members to the class declaration in `stream_reassembler.hh`

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

StreamReassembler::StreamReassembler(const size_t capacity) : _us_strs{}, _output(capacity), _capacity(capacity) {}

//! \details This function accepts a substring (aka a segment) of bytes,
//! possibly out-of-order, from the logical stream, and assembles any newly
//! contiguous substrings and writes them into the output stream in order.
void StreamReassembler::push_substring(const string &data, const size_t index, const bool eof) {
    if (eof) {
        _end_byte = index + data.size();
    }
    if (_output.bytes_written() >= _end_byte) {
        _output.end_input();
    }
    if (data.empty()) {
        return;
    }
    size_t us_cap = _output.remaining_capacity();
    if (index + data.size() <= _first_us_idx || index >= _first_us_idx + us_cap) {
        if (_output.bytes_written() >= _end_byte) {
            _output.end_input();
        }
        return;
    }
    // find newly covered ranges;
    // insert newly covered ranges with substring;
    // push from first unassembled and erase on the way;

    size_t l = max(index, _first_us_idx);
    size_t r = min(_first_us_idx + us_cap, index + data.size());
    segment seg = std::make_pair(l, l);

    vector<segment> segs_add;

    if (_us_strs.empty()) {
        segs_add.emplace_back(l, r);
    } else {
        auto l_iter = _us_strs.lower_bound(seg);
        if (l_iter == _us_strs.end()) {
            auto p_r = _us_strs.rbegin()->first.second;
            if (p_r < r) {
                segs_add.emplace_back(max(l, p_r), r);
            }
        }
        // filter out (l,r) inside existing segments
        else if (l_iter == _us_strs.begin() || std::prev(l_iter)->first.second < r) {
            size_t pre_r = l_iter == _us_strs.begin() ? l : max(l, std::prev(l_iter)->first.second);
            for (auto iter = l_iter; iter != _us_strs.end(); ++iter) {
                auto &cur_seg = iter->first;
                if (r <= cur_seg.first) {
                    segs_add.emplace_back(pre_r, r);
                    break;
                }
                if (r >= cur_seg.first) {
                    segs_add.emplace_back(pre_r, cur_seg.first);
                }
                if (r <= cur_seg.second) {
                    break;
                }
                pre_r = cur_seg.second;
            }
            auto p_l = _us_strs.rbegin()->first.second;
            if (p_l < r) {
                segs_add.emplace_back(p_l, r);
            }
        }
    }

    for (auto &&[sl, sr] : std::move(segs_add)) {
        _us_strs[make_pair(sl, sr)] = data.substr(sl - index, sr - sl);
        _us_bytes += sr - sl;
    }

    if (!_us_strs.empty()) {
        auto it = _us_strs.begin();
        while (it->first.first == _first_us_idx && it != _us_strs.end()) {
            size_t written = _output.write(std::move(it->second));
            _first_us_idx += written;
            _us_bytes -= written;
            ++it;
        }
        _us_strs.erase(_us_strs.begin(), it);
    }

    if (_output.bytes_written() >= _end_byte) {
        _output.end_input();
    }
}

size_t StreamReassembler::unassembled_bytes() const { return _us_bytes; }

bool StreamReassembler::empty() const { return unassembled_bytes() == 0; }
